package heap

import (
	"gitee.com/kklt1996/data-structure/array"
	"gitee.com/kklt1996/data-structure/util"
)

/*
	最大堆
	堆是一棵完全二叉树
*/
type MaxHeap struct {
	/*
		使用数组存储堆
	*/
	data *array.SliceArray

	/*
		使用comparator比较堆中元素大小,实现common.CompareAble接口和传入比较器形式二选一
		1 thisValue＞compareValue
		0 thisValue=compareValue
		-1 thisValue<compareValue
	*/
	comparator func(thisValue interface{}, compareValue interface{}) int
}

/*
	获取最大堆中元素的个数
*/
func (heap MaxHeap) GetSize() int {
	return heap.data.Size()
}

/*
	最大堆是否为空
*/
func (heap MaxHeap) IsEmpty() bool {
	return heap.data.IsEmpty()
}

/*
	O(log(n))
	向完全二叉树的数组表示中,添加元素
*/
func (heap MaxHeap) Add(value interface{}) {
	heap.data.AddLast(value)
	heap.siftUp(heap.data.Size() - 1)
}

func (heap MaxHeap) FindTop() (interface{}, error) {
	return heap.data.Get(0)
}

/*
	O(log(n))
	提取出最大的元素
*/
func (heap MaxHeap) ExtractTop() (interface{}, error) {
	// 获取最大的元素
	maxValue, err := heap.data.Get(0)
	if err != nil {
		return nil, err
	}
	// 将最后一个元素放置在第一个位置
	heap.data.Set(0, nil)
	lastValue, _ := heap.data.GetLast()
	heap.data.Set(0, lastValue)
	// 删除最后一个元素
	_, _ = heap.data.RemoveLast()
	// 如果元素个数大于1则需要siftDown
	if heap.data.Size() > 1 {
		heap.siftDown(0)
	}
	return maxValue, nil
}

/*
	取出最大的一个元素后,放入一个新的元素
*/
func (heap MaxHeap) Replace(value interface{}) (interface{}, error) {
	maxValue, err := heap.data.Get(0)
	if err != nil {
		return nil, err
	}
	// 替换堆顶的元素
	heap.data.Set(0, value)
	// 将对顶调整至合适的位置,使其满足最大堆
	heap.siftDown(0)
	return maxValue, nil
}

/*
	删除元素后将数组首部的元素调整至合适的位置
*/
func (heap MaxHeap) siftDown(index int) {
	// 从根节点开始,将左右孩子节点中比父亲节点大的节点中最大的节点和父亲节点交换位置,
	// 然后将交换的孩子节点作为父亲节点继续和左右孩子节点比较交换直到没有孩子节点大于父亲节点为止．
	size := heap.data.Size()
	indexValue, _ := heap.data.Get(index)
	for {
		leftChildIndex := heap.leftChildIndex(index)
		// 左节点的时候右节点也是为空的,右节点的索引比左节点的索引还要大1
		if leftChildIndex >= size {
			break
		}
		// 找出左右节点中最大的值,和最大的index
		var maxChildValue interface{}
		maxChildIndex := leftChildIndex
		maxChildValue, _ = heap.data.Get(leftChildIndex)
		rightChildIndex := heap.rightChildIndex(index)
		if rightChildIndex < size {
			rightChildValue, _ := heap.data.Get(rightChildIndex)
			if heap.compare(maxChildValue, rightChildValue) < 0 {
				maxChildValue = rightChildValue
				maxChildIndex++
			}
		}
		// 判断是否需要交换节点
		if heap.compare(maxChildValue, indexValue) > 0 {
			// 左右孩子节点中最大的节点大于父亲节点,需要交换孩子节点和父亲节点
			heap.data.Set(index, maxChildValue)
			heap.data.Set(maxChildIndex, indexValue)
			index = maxChildIndex
		} else {
			// 左右孩子节点都小于父亲节点,已经满足最大堆
			break
		}
	}
}

/*
	O(log(n))
	将新增的元素调整到完全二叉树中合适的位置
*/
func (heap MaxHeap) siftUp(index int) {
	// 将新增的元素和父亲节点对比,如果大于父亲节点,和父亲节点对调．然后继续和父亲节点比较,直到比父亲节点小或成为根节点为止
	value, _ := heap.data.Get(index)
	for index > 0 {
		parentIndex := heap.parentIndex(index)
		parentValue, _ := heap.data.Get(parentIndex)
		if heap.compare(value, parentValue) > 0 {
			// 将当前元素和根节点交换
			heap.data.Set(index, parentValue)
			heap.data.Set(parentIndex, value)
			// 更新当前节点的index
			index = parentIndex
		} else {
			// 比父亲节点小或者相等,终止siftUp
			break
		}
	}
}

func (heap MaxHeap) compare(thisValue interface{}, compareValue interface{}) int {
	if heap.comparator != nil {
		return heap.comparator(thisValue, compareValue)
	} else {
		return util.DefaultComparator(thisValue, compareValue)
	}
}

/*
	返回完全二叉树的数组表示中,一个索引所表示的元素的父亲节点的索引
*/
func (heap MaxHeap) parentIndex(index int) int {
	return (index - 1) / 2
}

/*
	返回完全二叉树的数组表示中,一个索引所表示的元素的左孩子节点的索引
*/
func (heap MaxHeap) leftChildIndex(index int) int {
	return index*2 + 1
}

/*
	返回完全二叉树的数组表示中,一个索引所表示的元素的右孩子节点的索引
*/
func (heap MaxHeap) rightChildIndex(index int) int {
	return index*2 + 2
}
